#!/usr/bin/env python
import socket
import time
from threading import Thread
from time import sleep
from re import search
import serial
import sys
import pty
import os
import subprocess
import traceback
import select
import pdb
import re
import argparse
import zmq
import json


#==============Serial Port====================
class SerialPortServer(object):
	"""Basic class for Serial Port.
	"""
	master_ = None
	slave_ = None
	slaveName_ = None
	serial_ = None

	def __init__(self, baudrate=9600):
		assert isinstance(baudrate,int),"'port' must be a integer."
		# Create a pty
		self.master_,self.slave_ = pty.openpty()
		self.slaveName_ = os.ttyname(self.slave_)
		print "slave device names:", self.slaveName_
		# Open pty
		print " Open the port to "+self.slaveName_
		sleep(1) 
		self.serial_=serial.Serial(self.slaveName_,baudrate,timeout=0.5) #pass parameter:slaveName
		print "Open port successful! the port information:"
		print self.serial_
		print "\n"

	def isOpen(self):
		return self.serial_.isOpen()

	def send(self,msg):
		msg = msg + "\r\n"
		ret = os.write(self.master_, msg)
		sys.stdout.flush()  #flush the buffer
		return ret

	def recv(self):
		return os.read(self.master_,1024)

	def close(self):
		os.close(self.master_)
		os.close(self.slave_)
		self.master_ = None
		self.slave_ = None


class MapTable(object):
	def __init__(self, map_table):
		self.map_table_ = map_table

	def getResponseString(self,input_str):
		# Input check
		if not isinstance(input_str,str):
			print("Error: 'input_str' must be a string.")
			raise TypeError("invalid type for input string.")

		for key in self.map_table_.keys():
			regex_ret = search(key,input_str)
			if regex_ret:
				print("Regex > {0}".format(key))
				print("Regex Ret > {0}".format(regex_ret.group(0)))
				ret_str = self.map_table_[key]
				return ret_str

		#pdb.set_trace()
		return self.map_table_["default"]


class SerialPortServerWrapper(Thread):
	"""A echo server for Serialport. It will be used to cheaten TM.
	"""
	def __init__(self,baudrate,map_table):
		super(SerialPortServerWrapper, self).__init__()
		self.server_ = SerialPortServer(baudrate=baudrate)
		self.map_table_ = MapTable(map_table = map_table)
		# Set thread as a Daemon thread. Once all of normal thread exit,
		# Daemon thread will crash instantly without any terminate process.
		# This will cause that some resources isn't released.
		self.setDaemon(True)	
	def run(self):
		try:
			while self.server_.isOpen():
				requ_str = self.server_.recv()
				if requ_str=="":
					break
				print("\r\n\r\nRecv from Serialport > {0}".format(requ_str))
				resp_str = self.map_table_.getResponseString(requ_str)
				print("Send to serialport > {0}".format(resp_str))
				self.server_.send(resp_str)
		finally:
			self.server_.close()


#=====================TCP Server & Client =============================
class Client(object):
	"""Basic class for TCP client connection.
	"""
	def __init__(self,sock,sock_from=None):
		self.source_flag_ = False
		if sock_from!=None :
			self.source_flag_ = True
			self.source_ip_,self.source_port_ = sock_from.getsockname()
			self.source_port_ = "{0}".format(self.source_port_)

		self.sock_ = sock
		self.sock_server = sock_from
		self.ip_,self.port_ = sock.getsockname()
		self.port_ = "{0}".format(self.port_)

	def getSock(self):
		return self.sock_

	def getIP(self):
		return self.ip_;

	def getPort(self):
		return self.port_

	def isTMConnection(self):
		return self.source_flag_

	def getSourceIP(self):
		return self.source_ip_

	def getSourcePort(self):
		return self.source_port_

	@staticmethod
	def getSockID(sock):
		return id(sock)

	def send(self, msg):
		#print("Send to {0}:{1} > {2}".format(self.ip_,self.port_,msg))
		self.sock_.send(msg)

	def recv(self,max_len=1024):
		return self.sock_.recv(max_len)


class Server(object):
	"""Basic class for TCP server.
	"""
	ctx = None
	def __init__(self, ip, port):
		assert isinstance(port,int),"'port' must be a integer."
		assert isinstance(ip,str),"'ip' must be a string."

		self.host_ = ip
		self.port_ = port
		self.sock_ = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		print("Server bind to {0}:{1}.".format(self.host_, self.port_))
		self.sock_.bind((self.host_, self.port_))
		self.sock_.listen(5)
		self.sock_.setblocking(False)		

	def getSock(self):
		return self.sock_

	@staticmethod
	def getSockID(sock):
		return id(sock)

	def accept(self):
		sock_conn, sock_addr = self.sock_.accept()
		sock_conn.setblocking(False)
		#print("Accept connection from {0}:{1}".format(sock_addr[0],sock_addr[1]))

		return Client(sock_conn,self.sock_)


class ClientConn(object):
	"""This is a client used to store connection information.
	"""
	def __init__(self,ip,port):
		# TODO
		assert isinstance(ip,str),"ip must be a string type."
		assert isinstance(port,int),"port must be a int type."
		addr = (ip,port)
		print("Client connect to {0}:{1}".format(ip,port))
		self.ip_ = ip
		self.port_ = port
		self.sock_ = socket.socket()
		self.sock_.connect(addr)
		self.sock_.setblocking(False)

	def getClient(self):
		return Client(self.sock_)

#====================Zmq publisher & subscriber================
class ZmqServer(object):
	def __init__(self, ctx,sock_type, ip, port):
		assert isinstance(ip,str),"'port' must be a string."
		self.ctx_ = ctx
		self.sock_ = self.ctx_.socket(sock_type)
		addr = "tcp://{0}:{1}".format(ip,port)
		print("Server connect to {0}.".format(addr))
		self.sock_.connect(addr)

	def send(self,msg):
		self.sock_.send(msg)

	def recv(self):
		return self.sock_.recv()


class ZmqClient(object):
	ctx = None
	def __init__(self,ctx,socktype,ip,port):
		self.ip_ = ip
		self.port_ = port
		self.ctx_ = ctx
		addr = "tcp://{0}:{1}".format(self.ip_,self.port_)
		self.sock_ = self.ctx_.socket(socktype)
		print("Client connect to {0}.".format(addr))
		self.sock_.connect(addr)

	def setFilter(self, filter=""):
		self.sock_.setsockopt(zmq.SUBSCRIBE,filter)

	def send(self, msg):
		self.sock_.send(msg)

	def recv(self):
		return self.sock_.recv()


class ZmqPublisher(ZmqServer):
	def __init__(self,ctx,ip,port):
		super(ZmqPublisher,self).__init__(ctx=ctx,sock_type=zmq.PUB,ip=ip,port=port)

	def recv(self):
		pass


class ZmqSubscriber(ZmqClient):
	def  __init__(self,ctx,ip,port):
		super(ZmqSubscriber,self).__init__(ctx=ctx,socktype=zmq.SUB,ip=ip,port=port)

	def send(self,msg):
		pass


#====================Public class=======================
class MessageParseException(Exception):
	def __init__(self,msg):
		self.msg_ = msg
	def getMsg(self):
		return self.msg_

class Message(object):
	def __init__(self):
		self.ip_ = None
		self.port_ = None
		self.data_ = None
		self.parse_flag_ = False

	def setFrom(self,ip,port,data):
		self.ip_ = ip
		self.port_ = port
		self.data_ = data
		self.parse_flag_ = True
		return self

	def setFromJSON(self,json_str):
		obj = json.loads(json_str)
		self.ip_ = str(obj.get('ip'))
		self.port_ = str(obj.get('port'))
		self.data_ = str(obj.get('data'))
		self.parse_flag_ = True
		return self

	def toJSON(self):
		obj = {'ip':self.ip_,'port':self.port_,'data':self.data_}
		return json.dumps(obj)

	def isSuccess(self):
		return self.parse_flag_


class ClientConnManager(object):
	"""There are problems still.
	"""
	def __init__(self):
		self.client_list_ = []

	def add(self,conn):
		assert isinstance(conn,Client),"Type must be Client."
		self.lock()
		self.client_list_.append(conn)
		self.unlock()

	def remove(self,conn):
		assert isinstance(conn,Client),"Type must be Client."
		self.lock()
		self.client_list_.remove(conn)
		self.unlock()

	def findConnByAddr(self,ip,port):
		"""Find connection by ip&port"""
		self.lock()
		if ip==None or port==None:
			return None
		# Search TM client connection
		for conn in self.client_list_:
			if conn.getIP()==ip and conn.getPort()==port:
				return conn
		self.unlock()
		return None

	def lock(self):
		pass

	def unlock(self):
		pass



class SlavePubHandler(Thread):
	"""Publish message to master. Send message is a JSON string.
	"""
	def computeObjID(self,sock):
		return id(sock)

	def buildSlaveServer_(self,slave_ip_list,slave_port_list):
		for ip in slave_ip_list:
			for port in slave_port_list:
				server = Server(ip,port)
				sn = self.computeObjID(server.getSock())
				# TODO FIX ME:
				self.inputs_.append(server.getSock())
				self.server_list_.append(server)
				self.conn_dic_[sn]  = server

	def __init__(self,ctx,master_ip,master_pport,slave_ip_list,slave_port_list):
		# TODO FIX ME: 
		super(SlavePubHandler, self).__init__()
		self.setDaemon(True)
		self.pub_ = ZmqPublisher(ctx,master_ip,master_pport)
		# List for 'select' listen
		self.inputs_ = []
		self.outputs_ = []
		# List for identify type.
		self.server_list_ = []
		self.client_manager_ = ClientConnManager()
		self.conn_dic_ = {}  # A dic for connection , <sn,connection>, sn = ip:port
		# Build slave Server
		self.buildSlaveServer_(slave_ip_list,slave_port_list)

	def getClientManager(self):
		return self.client_manager_

	def removeListenSock(self,conn):
		print("Clear TM-Connection from {0}:{1}.".format(conn.getIP(),conn.getPort()))
		self.inputs_.remove(conn.getSock())
		self.conn_dic_.pop(id(conn.getSock()))
		self.client_manager_.remove(conn)


	def readEvent(self,conn):
		"""Read message & Pack message & Send to corresponding conn."""
		msg_str = conn.recv(1024)
		if msg_str=="":
			# This pipe is broken,remove it
			self.removeListenSock(conn)
			return
		# pack message & Send to master.
		if self.pub_!=None:
			ip = conn.getIP()
			port = conn.getPort()
			data = msg_str
			msg = Message().setFrom(ip,port,data)
			print("Send to Master<{0}:{1}> {2}".format(msg.ip_,msg.port_,msg.toJSON()))
			self.pub_.send(msg.toJSON())
		else:
			print("This should never be arrived.")
			quit()

	def acceptEvent(self,server_conn):
		"""Accept connection"""
		# TODO record IP&Port.
		client_conn = server_conn.accept()

		sn = self.computeObjID(client_conn.getSock())
		self.inputs_.append(client_conn.getSock())
		self.conn_dic_[sn] = client_conn
		self.client_manager_.add(client_conn)
		print("Accept connection from {0}:{1}.".format(client_conn.getIP(),client_conn.getPort()))


	def run(self):
		while True:
			readable,writable,exceptional = select.select(self.inputs_,self.outputs_,self.inputs_)
			#pdb.set_trace()
			for i in readable:
				sn = self.computeObjID(i)
				conn = self.conn_dic_[sn]
				if conn in self.server_list_:
					"""conn is a server, accept"""
					self.acceptEvent(conn)
				else:
					"""conn is a client, read data"""
					self.readEvent(conn)

			for i in exceptional:
				print('handling exceptional condition for', i.getpeername()[0])
				sn = self.computeObjID(i)
				conn = self.conn_dic_[sn]
				self.inputs_.remove(conn.getSock())


class SlaveSubHandler(Thread):
	"""Subscribe message from master and transform it to TM. Recv message is
		a JSON string.
	"""
	def __init__(self,ctx,master_ip,master_sport,client_manager):
		super(SlaveSubHandler, self).__init__()
		self.setDaemon(True) 
		self.sub_ = ZmqSubscriber(ctx,master_ip,master_sport)
		self.sub_.setFilter()
		self.client_manager_ = client_manager


	def run(self):
		while True:
			try:
				# Recv message & Parse it
				json_str = self.sub_.recv()
				# pdb.set_trace()
				print("Recv from master> <{0}>".format(json_str))
				msg_obj = Message().setFromJSON(json_str)
				# Transform to TM
				ip = msg_obj.ip_
				port = msg_obj.port_
				data = msg_obj.data_
				send_conn = self.client_manager_.findConnByAddr(ip,port)
				if send_conn != None:
					print("--Transfer to TM("+ip+":"+port+") > "+data)
					send_conn.send(data)
				else:
					sn = ip+":"+port
					print("Failed to find target TM connection > "+sn)

			except Exception as e:
				traceback.print_exc()
				print("Abandon this message, continue to work.")


Blacktea_Map_Table = {
	"press":"START_OK\r",
	"release":"OK\r",
	"default":"ERR\r",
}


class DataCenter(object):
	def __init__(self,master_ip,master_pport,master_sport,slave_ip_list,slave_port_list):
		self.ctx_ = zmq.Context()
		self.blacktea_server_ = SerialPortServerWrapper(baudrate=115200,map_table=Blacktea_Map_Table)
		self.pub_handler_ = SlavePubHandler(self.ctx_,master_ip,master_pport,slave_ip_list,slave_port_list)
		self.sub_handler_ = SlaveSubHandler(self.ctx_,master_ip,master_sport,self.pub_handler_.getClientManager())

	def run(self):
		# Launch all of server
		self.blacktea_server_.start()  # Fixture control server
		# Send command to master.
		self.pub_handler_.start()
		# Subscribe command from master.
		self.sub_handler_.start()

		while True: sleep(1)


def createVirtualIP(ip_list):
	# TODO create a class
	for ip in ip_list:
		# FIX ME: check return value. Current method is invalid.
		ret = subprocess.call(["sudo","ifconfig","lo0","alias",ip])
		if ret!=0 :
			print("Failed to create Alias-IP. "+ip)
			raise Exception("Invalid alias IP. "+ip);


def releaseVirtualIP(ip_list):
	for ip in ip_list:
		# FIX ME: check return value. Current method is invalid.
		ret = subprocess.call(["sudo","ifconfig","lo0","-alias",ip])
		if ret!=0 :
			print("Failed to release Alias-IP. "+ip)


#================ Main ==========================
if __name__=="__main__":
	# Create two alias ip
	parser = argparse.ArgumentParser()
	parser.add_argument('-mip', '--master_ip', help='master ip.', type=str, default="127.0.0.1")
	parser.add_argument('-mpp', '--master_pub_port', help='port for publish message to master.', type=int, default=11001)
	parser.add_argument('-msp', '--master_sub_port', help='port for subscrib message from master.', type=int, default=11002)
	parser.add_argument('-ap', '--arm_port', help='port for arm board.', type=int, default=7600)
	parser.add_argument('-pp', '--power_port', help='port for power board.', type=int, default=8082)
	parser.add_argument('-dp', '--dut_port', help='port for dut board.', type=int, default=8081)

	parser.add_argument('-u1', '--uut1_ip', help='channel for dut 1.', type=str, default="169.254.1.30")
	parser.add_argument('-u2', '--uut2_ip', help='channel for dut 2.', type=str, default="169.254.1.31")

	args = parser.parse_args()
	master_ip = args.master_ip
	master_pport = args.master_pub_port;
	master_sport = args.master_sub_port;
	uut1 = args.uut1_ip
	uut2 = args.uut2_ip
	arm_port = args.arm_port
	power_port = args.power_port
	dut_port = args.dut_port

	ip_list = [uut1,uut2]
	port_list = [arm_port,power_port,dut_port]

	try:
		createVirtualIP(ip_list)

		data_center = DataCenter(master_ip=master_ip,master_pport=master_pport,master_sport=master_sport,slave_ip_list=ip_list,slave_port_list=port_list)
		data_center.run()
	except Exception as e:
		traceback.print_exc()
	finally:
		releaseVirtualIP(ip_list)
